import cv2
import joblib
import numpy as np
import os
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
from skimage.feature import hog


# 图像预处理函数：调整图像大小，灰度化，归一化
def preprocess_image(image, target_size=(128, 128)):
    image_resized = cv2.resize(image, target_size)
    gray_image = cv2.cvtColor(image_resized, cv2.COLOR_BGR2GRAY)
    normalized_image = gray_image / 255.0
    return normalized_image


# 加载数据并进行预处理
def load_and_preprocess_data(folder_path):
    images = []
    labels = []
    label_map = {}
    for label, plant_folder in enumerate(os.listdir(folder_path)):
        label_map[label] = plant_folder
        for image_name in os.listdir(os.path.join(folder_path, plant_folder)):
            image_path = os.path.join(folder_path, plant_folder, image_name)
            image = cv2.imread(image_path)
            preprocessed_image = preprocess_image(image)
            images.append(preprocessed_image)
            labels.append(label)
    return images, labels, label_map


# 提取图像的 HOG 特征
def extract_hog_features(images):
    hog_features = []
    for image in images:
        hog_feature = hog(image, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2), transform_sqrt=True)
        hog_features.append(hog_feature)
    return np.array(hog_features)


# 训练 SVM 模型
def train_svm_model(X_train, y_train):
    svm_model = SVC(kernel='linear')
    svm_model.fit(X_train, y_train)
    return svm_model


# 保存模型
def save_model(model, model_path):
    joblib.dump(model, model_path)

# 预测并评估模型准确率
def predict_and_evaluate(svm_model, X_test, y_test):
    y_pred = svm_model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    return accuracy


# 植物图像识别
def predict_plant(image_path, svm_model, label_map):
    image = cv2.imread(image_path)
    preprocessed_image = preprocess_image(image)
    hog_feature = extract_hog_features([preprocessed_image])
    prediction = svm_model.predict(hog_feature)
    plant_label = label_map[prediction[0]]
    return plant_label


# 主函数
def main():
    folder_path = "raw_train_image/"  # 包含每种植物照片的文件夹
    print("Loading data...")
    images, labels, label_map = load_and_preprocess_data(folder_path)
    print("Preprocessing data finished")

    print("extract_hog_features...")
    X = extract_hog_features(images)
    y = np.array(labels)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    print("Training SVM model...")
    svm_model = train_svm_model(X_train, y_train)
    accuracy = predict_and_evaluate(svm_model, X_test, y_test)
    print(f"Accuracy: {accuracy}")

    model_path = "model/svm_model.joblib"  # 模型路径
    save_model(svm_model, model_path)

    print("Predicting plant...")
    test_image_path = "img.png"  # 测试图像路径
    predicted_plant = predict_plant(test_image_path, svm_model, label_map)
    print(f"Predicted plant: {predicted_plant}")


if __name__ == "__main__":
    main()
